from enum import Enum, unique
import numpy as np
from numpy.typing import NDArray
import matplotlib.pyplot as plt
from typing import Optional

from common_util import check_is_3d_rotation_matrix, transform_3d_cloud
from common_util import euler_to_rotation_matrix, rotation_matrix_to_euler
from common_util import MATPLOTLIB_COLOR, init_axes, matplot_bounding_box, matplot_text
from common_util import bounding_box_to_points, transform_3d_cloud_to_image

__all__ = ['ObstacleCategory', 'PointXYZ', 'BoundingBox', 'Obstacle', 'MotionObstacle']

@unique
class ObstacleCategory(Enum):
    car = 'car'
    truck = 'truck'
    bus = 'bus'
    trailer = 'trailer'
    engineering_vehicle = 'engineering vehicle'
    bicycle = 'bicycle'
    motorcycle = 'motorcycle'
    pedestrian = 'pedestrian'
    traffic_cone = 'traffic cone'
    barrier = 'barrier'
    barrier_gate = 'barrier gate'
    debris = 'debris'
    ignore = 'ignore'

class PointXYZ():
    def __init__(self, x: float, y: float, z: float) -> None:
        self.x = float(x)
        self.y = float(y)
        self.z = float(z)

    def to_list(self) -> list:
        return [self.x, self.y, self.z]

    def to_dict(self) -> dict:
        return {'x': self.x, 'y': self.y, 'z': self.z}

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool=False) -> None:
        assert check_is_3d_rotation_matrix(rotation_matrix)
        point = transform_3d_cloud(np.array([self.x, self.y, self.z]).reshape((1, 3)),
                                   translation, rotation_matrix, inverse)
        self.x, self.y, self.z = point[0, 0], point[0, 1], point[0, 2]

class BoundingBox():
    def __init__(self, center, size, rotation) -> None:
        self.center = PointXYZ(*center)
        self.size = PointXYZ(*size)
        self.rotation = PointXYZ(*rotation)

    def to_dict(self) -> dict:
        return {
            'center': self.center.to_dict(),
            'size': self.size.to_dict(),
            'rotation': self.rotation.to_dict(),
        }

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool=False) -> None:
        assert check_is_3d_rotation_matrix(rotation_matrix)
        self.center.shift_coordinate_system(translation, rotation_matrix, inverse)
        heading_mat = euler_to_rotation_matrix([self.rotation.x, self.rotation.y, self.rotation.z], degree=False)
        if inverse:
            heading_mat = rotation_matrix.T.dot(heading_mat)
        else:
            heading_mat = rotation_matrix.dot(heading_mat)
        assert check_is_3d_rotation_matrix(heading_mat)
        self.rotation = PointXYZ(*rotation_matrix_to_euler(heading_mat))

    def matplot_on_bev(self, scale=0.3, rgb=MATPLOTLIB_COLOR['green'], alpha=1, axes: Optional[plt.Axes]=None) -> None:
        if axes is None:
            axes = init_axes()
        matplot_bounding_box(position=[self.center.x, self.center.y, self.center.z],
                             size=[self.size.x, self.size.y, self.size.z],
                             rotation=[self.rotation.x, self.rotation.y, self.rotation.z],
                             scale=scale, rgb=rgb, alpha=alpha, axes=axes)

    def matplot_on_camera(self, to_camera: NDArray, camera_intrinsic: NDArray, crop=(-1, -1),
                          scale=0.5, rgb=MATPLOTLIB_COLOR['green'], alpha=1, axes: Optional[plt.Axes]=None) -> None:
        box_pts = bounding_box_to_points(position=[self.center.x, self.center.y, self.center.z],
                                         size=[self.size.x, self.size.y, self.size.z],
                                         rotation=[self.rotation.x, self.rotation.y, self.rotation.z])
        box_pts = transform_3d_cloud_to_image(box_pts, to_camera, camera_intrinsic)
        assert box_pts.shape == (8, 3)
        if not np.all(box_pts[:, 2] > 0):
            return
        x_flt, y_flt = np.array([False]*8), np.array([False]*8)
        if crop[0] > 0:
            x_flt = (box_pts[:, 0] < 0) | (box_pts[:, 0] > crop[0])
        if crop[1] > 0:
            y_flt = (box_pts[:, 1] < 0) | (box_pts[:, 1] > crop[1])
        if np.all(x_flt | y_flt):
            return
        if axes is None:
            axes = init_axes()
            axes.set_xlabel('x axis')
            axes.set_ylabel('y axis')
            axes.axis('equal')
        axes.plot(box_pts[[0, 1, 2, 3, 0], 0], box_pts[[0, 1, 2, 3, 0], 1], c=rgb, lw=scale, alpha=alpha)
        axes.plot(box_pts[[4, 5, 6, 7, 4], 0], box_pts[[4, 5, 6, 7, 4], 1], c=rgb, lw=scale, alpha=alpha)
        for i in range(4):
            axes.plot(box_pts[[i, i+4], 0], box_pts[[i, i+4], 1], c=rgb, lw=scale, alpha=alpha)

class Obstacle():
    def __init__(self, id_, category: ObstacleCategory, center, size, rotation,
                 num_lidar_pts: int, num_radar_pts: int) -> None:
        assert isinstance(category, ObstacleCategory), f'ERROR: unsupported category: {category}'
        assert len(center) == len(size) == len(rotation) == 3, 'ERROR: bounding box info format wrong'
        assert num_lidar_pts >= 0 and num_radar_pts >= 0, 'ERROR: number of points should not be negative'
        self.id = id_
        self.category: ObstacleCategory = category
        self.box = BoundingBox(center, size, rotation)
        self.num_lidar_pts = num_lidar_pts
        self.num_radar_pts = num_radar_pts

    def to_dict(self) -> dict:
        self_dict = {
            'id': self.id,
            'category': self.category.value,
            'num_lidar_pts': self.num_lidar_pts,
            'num_radar_pts': self.num_radar_pts,
        }
        self_dict.update(self.box.to_dict())
        return self_dict

    @classmethod
    def from_dict(cls, data: dict) -> 'Obstacle':
        return cls(id_=data['id'],
                   category=ObstacleCategory(data['category']),
                   center=[data['center']['x'], data['center']['y'], data['center']['z']],
                   size=[data['size']['x'], data['size']['y'], data['size']['z']],
                   rotation=[data['rotation']['x'], data['rotation']['y'], data['rotation']['z']],
                   num_lidar_pts=data['num_lidar_pts'],
                   num_radar_pts=data['num_radar_pts'])

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool=False) -> None:
        self.box.shift_coordinate_system(translation, rotation_matrix, inverse)

    def matplot_on_bev(self, scale=0.3, rgb=MATPLOTLIB_COLOR['green'], alpha=1, axes: Optional[plt.Axes]=None) -> None:
        if axes is None:
            axes = init_axes()
        self.box.matplot_on_bev(scale, rgb, alpha, axes)
        matplot_text(self.box.center.x, self.box.center.y, self.category.value,
                     scale=2, rgb=rgb, alpha=alpha, axes=axes)

    def matplot_on_camera(self, to_camera: NDArray, camera_intrinsic: NDArray, crop=(-1, -1),
                          scale=0.5, rgb=MATPLOTLIB_COLOR['green'], alpha=1, axes: Optional[plt.Axes]=None) -> None:
        if axes is None:
            axes = init_axes()
            axes.set_xlabel('x axis')
            axes.set_ylabel('y axis')
            axes.axis('equal')
        self.box.matplot_on_camera(to_camera, camera_intrinsic, crop, scale, rgb, alpha, axes)

class MotionObstacle(Obstacle):
    def __init__(self, id_, category, center, size, rotation, velocity, acceleration,
                 num_lidar_pts: int, num_radar_pts: int) -> None:
        super().__init__(id_, category, center, size, rotation, num_lidar_pts, num_radar_pts)
        assert len(velocity) == len(acceleration) == 3
        self.velocity = PointXYZ(*velocity)
        self.acceleration = PointXYZ(*acceleration)

    def to_dict(self) -> dict:
        self_dict = super().to_dict()
        self_dict.update({
            'velocity': self.velocity.to_dict(),
            'acceleration': self.acceleration.to_dict(),
        })
        return self_dict

    @classmethod
    def from_dict(cls, data: dict) -> 'MotionObstacle':
        return cls(id_=data['id'],
                   category=ObstacleCategory(data['category']),
                   center=[data['center']['x'], data['center']['y'], data['center']['z']],
                   size=[data['size']['x'], data['size']['y'], data['size']['z']],
                   rotation=[data['rotation']['x'], data['rotation']['y'], data['rotation']['z']],
                   velocity=[data['velocity']['x'], data['velocity']['y'], data['velocity']['z']],
                   acceleration=[data['acceleration']['x'], data['acceleration']['y'], data['acceleration']['z']],
                   num_lidar_pts=data['num_lidar_pts'],
                   num_radar_pts=data['num_radar_pts'])

    def shift_coordinate_system(self, translation: NDArray,
                                rotation_matrix: NDArray, inverse: bool = False) -> None:
        super().shift_coordinate_system(translation, rotation_matrix, inverse)
        self.velocity = PointXYZ(*transform_3d_cloud(
            np.array([self.velocity.x, self.velocity.y, self.velocity.z]).reshape((1, 3)),
            np.zeros((1, 3)), rotation_matrix, inverse).reshape(3))
        self.acceleration = PointXYZ(*transform_3d_cloud(
            np.array([self.acceleration.x, self.acceleration.y, self.acceleration.z]).reshape((1, 3)),
            np.zeros((1, 3)), rotation_matrix, inverse).reshape(3))
